import json
import av
import cv2
import socket
import struct
import time
import threading
from logger import log
from settings import settings
import numpy as np

streams = set()
client_socket = None
heartbeat_interval = 30  # 心跳间隔时间（秒）


def stream_video(rtsp_url, id):
    log.info(f"Starting stream for id {id} from {rtsp_url}")
    cap = cv2.VideoCapture(rtsp_url)  # 打开摄像头
    while cap.isOpened() and id in streams:
        ret, frame = cap.read()
        if not ret:
            continue

        img = cv2.resize(frame, (640, 360))
        _, buffer = cv2.imencode('.jpg', img)
        frame_data = buffer.tobytes()

        try:
            # 发送帧大小
            header = struct.pack("QQ", id, len(frame_data))
            client_socket.sendall(b'@#@#' + header + frame_data)
        except Exception as e:
            log.error(f"Error sending frame for id {id}: {e}")
            break  # 出现错误时退出循环

        time.sleep(0.1)  # 控制帧率

    cap.release()
    log.info(f"Stopped stream for id {id}")


def stream_video_av(rtsp_url, id):
    log.info(f"stream_video_av:{id},{rtsp_url}")
    frame_count = 0  # 帧计数器
    try:
        container = av.open(rtsp_url,
                            format='rtsp',
                            options={'rtsp_transport': 'tcp',
                                     'framerate': '20',
                                     'video_size': '640x360'},
                            timeout=25)
    except Exception as e:
        log.warn(f"generate_frames_av Failed to open RTSP stream: {e}")
        # 创建一个固定图片
        fixed_image = np.zeros((360, 640, 3), dtype=np.uint8)  # 黑色图片
        cv2.putText(fixed_image, 'Failed to open stream', (10, 180),
                    cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 255, 255), 2, cv2.LINE_AA)
        ret, buffer = cv2.imencode('.jpg', fixed_image)
        if ret:
            frame_data = buffer.tobytes()
            # 发送帧大小
            header = struct.pack("QQ", id, len(frame_data))
            # log.info(f"header:{header},{len(header)}")
            client_socket.sendall(b'@#@#' + header + frame_data)
        return None
    for frame in container.decode(video=0):
        if id not in streams:
            break
        frame_count += 1
        # 每读取到五帧只处理一帧
        if frame_count % 5 == 0:
            # 将帧转换为numpy数组
            img = frame.to_ndarray(format='bgr24')
            # 将帧调整为640x360
            img = cv2.resize(img, (640, 360))
            # 将帧编码为JPEG格式
            ret, buffer = cv2.imencode('.jpg', img)
            if not ret:
                continue
            frame_data = buffer.tobytes()
            # 发送帧大小
            header = struct.pack("QQ", id, len(frame_data))
            # log.info(f"header:{header},{len(header)}")
            client_socket.sendall(b'@#@#' + header + frame_data)


def send_heartbeat():
    global client_socket
    while True:
        time.sleep(heartbeat_interval)
        try:
            log.info("send Heartbeat")
            if client_socket:
                heartbeat_header = struct.pack("QQ", 0, 0)
                client_socket.sendall(b'@#@#' + heartbeat_header)
            else:
                log.warn("Heartbeat exit")
        except Exception as e:
            log.error(f"Heartbeat error: {e}")


def handle_server_connection(server_address, edge_code):
    global client_socket
    while True:
        try:
            client_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            client_socket.settimeout(60)  # 设置接收超时为60秒
            client_socket.connect(server_address)
            log.info("服务器连接成功")
            ss = []
            for s in settings.streams:
                ss.append({"id": s["id"], "name": s["name"]})
            init_data = json.dumps({"edge_code": edge_code, "streams": ss})
            json_data = init_data.encode()
            data_length = struct.pack('!I', len(json_data))
            client_socket.sendall(data_length + json_data)

            while True:
                try:
                    # 读取4字节命令长度
                    length_data = recvall(client_socket, 4)
                    if not length_data:
                        raise ConnectionResetError("Server closed connection")
                    cmd_length = struct.unpack('!I', length_data)[0]

                    # 读取完整的JSON命令
                    cmd_data = recvall(client_socket, cmd_length)
                    if not cmd_data:
                        raise ConnectionResetError("Server closed connection")
                    cmd_str = cmd_data.decode()
                    log.info(f"cmd_str:{cmd_str}")
                    obj = json.loads(cmd_str)
                    if "start" == obj['cmd']:
                        streams.add(obj['id'])
                        for stream in settings.streams:
                            if stream["id"] == obj['id']:
                                stream_url = stream["rtsp"]
                                if stream['cameraRtsp'] is not None and stream['cameraRtsp'] != '':
                                    stream_url = stream['cameraRtsp']
                                threading.Thread(target=stream_video_av, args=(stream_url, obj['id'],),
                                                 daemon=True).start()
                    elif "stop" == obj['cmd']:
                        streams.discard(obj['id'])
                    elif "heartbeat" == obj['cmd']:
                        log.info("heartbeat")
                except socket.timeout:
                    log.warning("client_socket timeout")
                    break  # 超时则继续等待新命令
                except json.JSONDecodeError as e:
                    log.error(f"JSON decode error: {e}")
                except struct.error as e:
                    log.error(f"Struct unpack error: {e}")
        except Exception as e:
            log.exception(e)
            log.error(f"Connection error: {e}")
            log.info(f"Reconnecting in 10 seconds...")
            time.sleep(10)


def recvall(sock, n):
    """Helper function to receive n bytes or return None if EOF is hit"""
    data = b''
    while len(data) < n:
        packet = sock.recv(n - len(data))
        if not packet:
            return None
        data += packet
    return data


def start():
    log.info(f"vs_host:{settings.vs_host}")
    server_address = (settings.vs_host, 9999)
    threading.Thread(target=handle_server_connection, args=(server_address, settings.code), daemon=True).start()
    threading.Thread(target=send_heartbeat, daemon=True).start()


if __name__ == '__main__':
    edge_code = 'lipeng_test'
    # 配置初始化，根据需要调整
    # settings.init(edge_code, "platform.sxgjwl.com", None, "video.sxgjwl.com")
    # settings.init(edge_code, "platform.sxgjwl.com", None, "127.0.0.1")
    settings.init(edge_code, "platform.wechain360.com", None, "127.0.0.1")
    start()
    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        log.info("客户端退出")
